# ----------------------------------------------------------------------------------
# Microsoft Developer & Platform Evangelism
# 
# Copyright (c) Microsoft Corporation. All rights reserved.
# 
# THIS CODE AND INFORMATION ARE PROVIDED "AS IS" WITHOUT WARRANTY OF ANY KIND, 
# EITHER EXPRESSED OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE IMPLIED WARRANTIES 
# OF MERCHANTABILITY AND/OR FITNESS FOR A PARTICULAR PURPOSE.
# ----------------------------------------------------------------------------------
# The example companies, organizations, products, domain names,
# e-mail addresses, logos, people, places, and events depicted
# herein are fictitious.  No association with any real company,
# organization, product, domain name, email address, logo, person,
# places, or events is intended or should be inferred.
# ----------------------------------------------------------------------------------
#
# THIS IS SAMPLE CODE. BUGS ARE EXPECTED. SECURITY IS NOT.
#
# This module creates the supporting structures (queues, tables) and defines the
# functions to submit jobs, retrieve status, etc...
# It also shows examples of use in the test section.
#

#name of input queue
qname = 'inqueue'
#name of job status table
jobtable = 'jobtable'
#storage account
storage_account='your storage account'
storage_key='your storage key'

#name of command queue (actually a service bus topic)
cqname = 'cqueue'
#name of result queue (Running, Completed, Failed, actually a service bus topic)
rqname = 'outqueue'
#service bus namespace
sbname='gerespy'
#service bus key
sbkey='your service bus key'
#service bus issuer
sbissuer='owner'


from datetime import *
from subprocess import *
from ast import *
from azure.storage import *
from azure.servicebus import *
from socket import *
from time import *
from azure import *
from os import *
from string import *
from random import *
#from threading import *
#some ui work
#from Tkinter import *

#open connection to azure storage queue and table and service bus
queue_service = QueueService(account_name=storage_account, account_key=storage_key)
table_service = TableService(account_name=storage_account, account_key=storage_key)
bus_service = ServiceBusService(service_namespace=sbname, account_key=sbkey, issuer=sbissuer)


#create the command and result status topics
def sbsetup(sbqname):
    try:
        result=bus_service.get_topic(sbqname)
    except WindowsAzureError as error:
        print 'Service bus topic ',sbqname,' does not exist. Creating...'
        bus_service.create_topic(sbqname)


#create the job queue
def qsetup(qname):
    try:
        result=queue_service.get_queue_metadata(qname)
    except WindowsAzureMissingResourceError as error:
        print 'Queue ', qname, 'does not exist. Creating...'
        queue_service.create_queue(qname)


#create the status table
def tablesetup(jobtable):
    try:
        result=table_service.query_tables(jobtable)
    except WindowsAzureMissingResourceError as error:
        print 'Table ', jobtable, 'does not exist. Creating...'
        table_service.create_table(jobtable)

#submit a job to the queue
def qsub(qname, jobtable, command, datain, dataout):
    #arbitrarily chose 16 bytes as random jobid length. change to whatever suits.
    jobid="".join([choice(letters+digits) for x in range(1, 16)])
    jobobject = {"JobID": jobid, "TimeIn": datetime.now().isoformat(), "Command": command, "DataIn": datain, "DataOut": dataout}
    print "Submitting: ", jobobject
    queue_service.put_message(qname, jobobject)
    #update status table to submitted
    #get the node name, which will be used as partition key for status table
    node=gethostname()
    jobstatus = {"Status": 'Submitted', "TimeIn": datetime.now().isoformat(), "TimeUpdated": datetime.now().isoformat()}
    table_service.insert_or_replace_entity(table_name=jobtable, partition_key=node, row_key=jobid, entity=jobstatus)
    #generate a notification of status
    msg=Message('Submitted',custom_properties={'JobID': jobid})
    bus_service.send_topic_message(rqname,msg)
    print "Status updated: ", jobstatus
    return jobid

#update job status
def update_status (table, node, id, newstatus):
        jobstatus=table_service.get_entity(table, node, id)
        print 'Current status: ', jobstatus.Status
        jobstatus.Status=newstatus
        jobstatus.TimeUpdated=datetime.now().isoformat()
        table_service.insert_or_replace_entity(table,node,id,jobstatus)
        msg=Message(newstatus,custom_properties={'JobID': id})
        bus_service.send_topic_message(rqname,msg)
        print 'New status: ', jobstatus.Status

#list all jobs in a queue
def qlist(qname):
    qmetadata = queue_service.get_queue_metadata(qname)
    nofmsg = qmetadata['x-ms-approximate-messages-count']
    print 'Messages in queue: ', nofmsg
    if nofmsg > 32:
        print 'Max supported messages for peek operation is 32. Only top 32 jobs will be listed.'
        nofmsg=32
    messages = queue_service.peek_messages(qname,nofmsg)
    print '%-16s\t%-26s\t%16s\t%16s\t%16s' % ('JobID','TimeIn', 'Command', 'DataIn', 'DataOut')
    for message in messages:
        jobin = literal_eval(message.message_text)
        print '%-16s\t%-26s\t%16s\t%16s\t%16s' % (jobin['JobID'],jobin['TimeIn'],jobin['Command'],jobin['DataIn'],jobin['DataOut'])

#list status of jobs
def joblist(jobtable):
    tasks=table_service.query_entities(jobtable)
    print '%-10s\t%-16s\t%-10s\t%-26s\t%-26s' %('Node','JobID','Status','TimeIn','TimeUpdated')
    for task in tasks:
        print '%-10s\t%-16s\t%-10s\t%-26s\t%-26s' % (task.PartitionKey,task.RowKey,task.Status,task.TimeIn,task.TimeUpdated)
        
#query status of 1 job
def jobquery(jobtable,jobid):
    filter="RowKey eq '"+jobid+"'"
    print filter
    tasks=table_service.query_entities(jobtable,filter)
    print '%-10s\t%-16s\t%-10s\t%-26s\t%-26s' %('Node','JobID','Status','TimeIn','TimeUpdated')
    for task in tasks:
        print '%-10s\t%-16s\t%-10s\t%-26s\t%-26s' % (task.PartitionKey,task.RowKey,task.Status,task.TimeIn,task.TimeUpdated)
    return task.Status

#query status of jobs on 1 node
def nodequery(jobtable,nodeid):
    filter="PartitionKey eq '"+nodeid+"'"
    print filter
    tasks=table_service.query_entities(jobtable,filter)
    print '%-10s\t%-16s\t%-10s\t%-26s\t%-26s' %('Node','JobID','Status','TimeIn','TimeUpdated')
    for task in tasks:
        print '%-10s\t%-16s\t%-10s\t%-26s\t%-26s' % (task.PartitionKey,task.RowKey,task.Status,task.TimeIn,task.TimeUpdated)

def jobcancel(jobtable,jobid):
    try:
        filter="RowKey eq '"+jobid+"'"
        print filter
        tasks=table_service.query_entities(jobtable,filter)
        print '%-10s\t%-16s\t%-10s\t%-26s\t%-26s' %('Node','JobID','Status','TimeIn','TimeUpdated')
        for task in tasks:
            print '%-10s\t%-16s\t%-10s\t%-26s\t%-26s' % (task.PartitionKey,task.RowKey,task.Status,task.TimeIn,task.TimeUpdated)
            if task.Status=='Running':
                msg2=Message('CANCEL',custom_properties={'JobID': jobid})
                bus_service.send_topic_message(cqname,msg2)
                print 'Job ',jobid,' executing. Cancellation request sent.'
            else:
                jobstatus = {"Status": 'Cancelled', "TimeIn": task.TimeIn, "TimeUpdated": datetime.now().isoformat()}
                table_service.insert_or_replace_entity(table_name=jobtable, partition_key=task.PartitionKey, row_key=jobid, entity=jobstatus)
                msg=Message('Cancelled',custom_properties={'JobID': jobid})
                bus_service.send_topic_message(rqname,msg)
                print 'Job ',jobid,' already dealt with.'
    except WindowsAzureError as error:
        print 'The jobid ',jobid,' cannot be found.'
    

    

###Setup part###
qsetup(qname)
tablesetup(jobtable)
sbsetup(cqname)
sbsetup(rqname)
bus_service.create_subscription(rqname, 'AllResults')

###Test Part###

#thread to read notifications asynchronously
#this had to be put in a separate python process because the Message directive in servicebus conflicts with the 
#message directive in TK!!! yet another bug!!
#class Notifier(Thread):
#    def run(self):
#        root=Tkinter.Tk()
#        root.title('Status Notifications')
#        msgtext=StringVar()
#        w=Tkinter.Label(root, textvariable=msgtext, width=50)
#        w.pack()
#        msgtext.set("Job status notifications will appear here.")
#        w.update()
#        while True:
#            try:
#                msg=bus_service.receive_subscription_message(rqname,'AllResults')
#            #print msg.custom_properties
#            #m=literal_eval(msg.custom_properties)
#                m=str(msg.custom_properties['jobid'])
#                msgtext.set('JobID: '+m+' '+msg.body)
#            #w.mainloop()
#            #print msg.custom_properties,' ',msg.body
#            except WindowsAzureError as error:
#                msgtext.set('No more notifications. Waiting...')    
#            root.update_idletasks()
#            sleep(5)
            


#populate queue with job objects
qn = raw_input('N. of elements in queue? ')
print 'Submitting', qn,' tasks to the storage queue...'
index = 1
#c:\pi\pi.exe is supposed to exist on the compute nodes. How to get it there is out of scope for this project.
#We simply added it to the VM template.
#pi.exe <n. of points> estimates pi with Laplace's statistical method. It is a good example of parametric sweep.
while index <= int(qn):
    j=qsub(qname, jobtable, '\\Pi\\pi.exe', str(index*100000), '\\pi\\pi.out')
    index = index + 1

raw_input('Press any key to continue.')


print 'Invoking the query functions...'

print 'Listing content of job queue...'
qlist(qname)
raw_input('Press any key to continue.')
print 'Listing content of job table...'
joblist(jobtable)
raw_input('Press any key to continue.')
print 'Querying status of specific job...'
jobquery(jobtable,j)
raw_input('Press any key to continue.')
print 'Listing jobs on a specific node...'
node=gethostname()
nodequery(jobtable,node)

raw_input('Press any key to cancel last job.')

#cancel last job to test command
print 'JobID: ',j
jobcancel(jobtable,j)
sleep(5)
#read status notifications from service bus topic
#print 'Reading notifications from service bus'
#NT=Notifier()
#NT.setDaemon(True)
#NT.start()
print 'Querying status of specific job...'
jobquery(jobtable,j)
raw_input('Press any key to finish.')
